"""
@Time    : 2019/10/7 9:57
@Author  : CcH
"""
from tqdm import tqdm
import torch
class Text_DataLoader():
    def __init__(self,path):
        self.path = path
        self.data = self.read_data()

    def read_data(self):
        """key,value"""
        data = {}
        with open(self.path, 'r', encoding="utf-8") as fr:
            for line in tqdm(fr.readlines()):
                id_text = line.strip().split('\t')
                data[id_text[0]] = id_text[1]
        return data

    def get_word2idx(self):
        word2idx = {}
        idx = 0
        for value in self.data.values():
            if value not in word2idx:
                word2idx[value] = idx
                idx += 1
        return word2idx

    def get_train_labels_info(self):
        word2idx = self.get_word2idx()
        train_labels = {}
        for key,value in self.data.items():
            train_labels[key] = word2idx[value]
        return train_labels

    def sava_data(self):
        word2idx = self.get_word2idx()
        train_labels = self.get_train_labels_info()
        print()
        data = {
            "word2idx": word2idx,
            "train_labels": train_labels
        }
        torch.save(data, "train_labels.pt")






if __name__ == '__main__':
    path = r"D:\XW_Bank\LipRecognition\train\lip_train.txt"
    dataloader = Text_DataLoader(path)
    dataloader.sava_data()
    data = torch.load("train_labels.pt")
    print(data["train_labels"])
    print(len(data["word2idx"]))
